import platgo as pg
import numpy as np


"""
------------------------------- Reference --------------------------------
 R. Eberhart and J. Kennedy, A new optimizer using particle swarm theory,
 Proceedings of the 6th International Symposium on Micro Machine and Human
 Science, 1995, 39-43.
"""


class PSO(pg.Algorithm):

    type: dict = {'single': True, 'multi': False, 'many': False, 'real': True, 'binary': False, 'permutation': False,
                  "large": True, 'expensive': False, 'constrained': False, 'preference': False, 'multimodal': False,
                  'sparse': False, 'gradient': False}

    def __init__(self, maxgen: int = None, problem: pg.Problem = None, weight: float = 0.4):
        # TODO 未考虑约束
        super().__init__(maxgen=maxgen, problem=problem)
        self.name = "PSO"
        self.weight = weight

    def go(self, N: int = None, population: pg.Population = None) -> pg.Population:
        """
        if population is None, generate a new population with population size
        :param N: population size
        :param population: population to be optimized
        :return: pop
        """
        # TODO 未考虑约束，应根据约束更新种群
        assert N or population, "N and population can't be both None"
        if population is None:
            pop = self.problem.init_pop(N)
        else:
            pop = population
            self.problem.N = pop.decs.shape[0]

        self.problem.cal_obj(pop)
        # 创建pbest副本
        pbest = pop.copy()
        best = np.argmin(pbest.objv)
        gbest = pbest.decs[best]

        while self.not_terminal(pop):
            pop = pg.operators.Operator_PSO(pop, pbest, gbest, self.weight)
            self.problem.cal_obj(pop)
            replace = (pbest.objv > pop.objv).flatten()
            pbest[replace] = pop[replace]
            best = np.argmin(pbest.objv)
            gbest = pbest.decs[best]
        return pop
